package ai.mspbots.poc.es.service;

import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Service
@Slf4j
public class ElasticsearchSaveService {

    @Autowired
    private RestHighLevelClient client;

    public ElasticsearchSaveService(RestHighLevelClient client) {
        this.client = client;
    }

    /**
     * Save data to Elasticsearch
     * @param index 索引
     * @param id 文档id
     * @param name report name
     * @param description report description
     * @param nameVector report name vector
     * @param descriptionVector report description vector
     */
    public void saveToElasticsearch(String index, String id, String name, String description, INDArray nameVector, INDArray descriptionVector) {
        double[] nameVectorValue;

        try {
            nameVectorValue = nameVector.toDoubleVector();
        } catch (Exception e) {
            nameVectorValue = null;
        }

        double[] descriptionVectorValue;
        try {
            descriptionVectorValue = descriptionVector.toDoubleVector();
        } catch (Exception e) {
            descriptionVectorValue = null;
        }

        Map<String, Object> document = new HashMap<>();
        document.put("id", id);
        document.put("name", name);
        document.put("name_vector", nameVectorValue);
        document.put("description", description);
        document.put("description_vector", descriptionVectorValue);
        // List<String> suggestList = new ArrayList<>();
        // suggestList.add(name);
        // document.put("suggest", suggestList);
        document.put("weight", 0);
        IndexRequest request = new IndexRequest(index)
                .id(id)
                .source(document);
        try {
            client.index(request, RequestOptions.DEFAULT);
        } catch (IOException e) {
            log.error("Failed to save data to Elasticsearch: {}", e.getMessage(), e);
        }
    }
}